package net.javlov.vf.fa;

import java.util.Arrays;

import net.javlov.Option;
import net.javlov.vf.OptionValues;

public class BasicLinearFATraces implements LinearFATraces {

	protected double traces[][], threshold, gamma, lambda;
	protected int nonZero[][], numNonZero[], defaultFeatIdx[], numOpt, numFeats;
	protected boolean replacing = true;
	
	public BasicLinearFATraces(double gamma, double lambda, int numOpt, int numFeats) {
		this.gamma = gamma;
		this.lambda = lambda;
		this.numOpt = numOpt;
		this.numFeats = numFeats;
		init();
		setThreshold(0.01);
	}
	
	public void init() {
		traces = new double[numOpt][numFeats];
		nonZero = new int[numOpt][(int)Math.min(numFeats, 5000)]; //assuming lambda always < 1
		numNonZero = new int[numOpt];
	}
	
	@Override
	public void update(double[] stateCoding, Option<?> o, OptionValues values) {
		if ( defaultFeatIdx == null )
			makeDefaultIdx(stateCoding);
		update(stateCoding, o, defaultFeatIdx, values);
	}
	
	protected void makeDefaultIdx(double[] stateCoding) {
		defaultFeatIdx = new int[stateCoding.length];
		for (int i = 0; i < stateCoding.length; i++) {
			defaultFeatIdx[i] = i;
		}
	}
	
	@Override
	public void update(double[] featVals, Option<?> opt, int[] featIdx, OptionValues values) {
		for ( int o = 0; o < traces.length; o++ ) {
			for ( int f = numNonZero[o]-1; f >= 0; f-- ) {
				decay(values, opt, o, nonZero[o][f]);

				//remove trace if below threshold
				if ( traces[o][nonZero[o][f]] < threshold ) {
					traces[o][nonZero[o][f]] = 0;
					nonZero[o][f] = nonZero[o][numNonZero[o]-1]; //copy index from head to current loc
					numNonZero[o]--;
				}
			}
		}
		
		int optId = opt.getID();
		for ( int f = 0; f < featIdx.length; f++ ) {
			if ( traces[optId][featIdx[f]] < threshold ) //only increase nonzero count if trace was 0 before
				nonZero[optId][numNonZero[optId]++] = featIdx[f];
			visit(optId, featVals[f], featIdx[f]);
		}
	}
	
	@Override
	public double[][] get() {
		return traces;
	}

	@Override
	public int[][] getLiveIndices() {
		return nonZero;
	}
	
	@Override
	public int[] getNumNonZero() {
		return numNonZero;
	}
	
	/**
	 * "Visit" option optId for feature index featIdx.
	 * 
	 * This default implementation sets the trace for (featIdx,optId) to featVal
	 * for replacing traces, and to e(featIdx,optId)+featVal for accumulating.
	 * 
	 * For example, for binary features, this means that e = 1 for replacing traces,
	 * and e = e + 1 for accumulating ones.
	 * 
	 * Override if anything more complicated should happen.
	 * 
	 * @param optId the option index
	 * @param featVal the feature value
	 * @param featIdx the feature index
	 */
	public void visit(int optId, double featVal, int featIdx) {
		//TODO questionable if this should be done this way for non-binary features
		if ( replacing )
			traces[optId][featIdx] = featVal;
		else
			traces[optId][featIdx] += featVal;
	}
	
	/**
	 * Decay the trace for feature featIdx and option opt, and last visited state-option pair (s,o).
	 * 
	 * This default implementation simply always decays by gamma*lambda. Override if anything
	 * more complicated should happen.
	 * 
	 * @param s the state last visited by the agent
	 * @param o the option last picked by the agent
	 * @param opt the option index for which to update the trace
	 * @param featIdx the feature index for which to update the trace
	 */
	public void decay(OptionValues vals, Option o, int opt, int featIdx) {
		traces[opt][featIdx] *= gamma*lambda;
	}
	
	public void setReplacing(boolean r) {
		replacing = r;
	}
	
	public boolean isReplacing() {
		return replacing;
	}

	public double getThreshold() {
		return threshold;
	}

	public void setThreshold(double threshold) {
		this.threshold = threshold;
	}

}
